package org.dromara.web.plugins.provider.imageImpl;


import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.dromara.common.core.exception.ServiceException;
import org.dromara.common.mongo.ChatMessageMongoDao;
import org.dromara.common.mongo.SessionMongoDao;
import org.dromara.dao.business.domain.Model;
import org.dromara.dao.business.domain.ModelParam;
import org.dromara.web.plugins.provider.TextToImageProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;

@Slf4j
@Component
public class SDXLV1ImageProvider implements TextToImageProvider {


    public static final String MODEL_TYPE = "sdxlv10";

    private static final String URL = "https://api.stability.ai/v1/generation/stable-diffusion-v1-6/text-to-image";
    private static final String API_KEY = "sk-l8ZwwA9Rs6luswXI2V4o4OzrH8IzCMQN1cwJuEIjLib39bOy";

    @Autowired
    OkHttpClient okHttpClient;
    @Autowired
    ThreadPoolTaskExecutor threadPoolTaskExecutor;
    @Autowired
    SessionMongoDao sessionMongoDao;
    @Autowired
    ChatMessageMongoDao chatMessageMongoDao;

    public static final Map<String, String> sizeMap = new HashMap<>();
    static {
        sizeMap.put("9:21", "576x1344");
        sizeMap.put("9:16", "576x1024");
        sizeMap.put("2:3", "1024x1536");
        sizeMap.put("4:5", "1024x1280");
        sizeMap.put("1:1", "1024x1024");
        sizeMap.put("5:4", "1280x1024");
        sizeMap.put("3:2", "1536x1024");
        sizeMap.put("16:9", "1024x576");
        sizeMap.put("21:9", "1344x576");
    }


    @Override
    public String generateImage(Model model, String prompt, ModelParam modelParam, String scale) {

        MultipartBody.Builder builder = new MultipartBody.Builder().setType(MultipartBody.FORM);
        builder.addFormDataPart("text_prompts", prompt);
        builder.addFormDataPart("output_format", "jpeg");
        if (modelParam != null) {
            builder.addFormDataPart("style_preset", modelParam.getPromt());
        }
        builder.addFormDataPart("aspect_ratio", scale);

        String width = "512";
        String height = "512";
        String pixels = sizeMap.get(scale);
        if (pixels != null) {
            String[] split = pixels.split("x");
            width = split[0];
            height = split[1];
        }
        builder.addFormDataPart("width", width);
        builder.addFormDataPart("height", height);


        RequestBody requestBody = builder.build();

        Request request = new Request.Builder()
            .url(URL)
            .header("Authorization", "Bearer " + API_KEY)
            .header("Accept", "application/json")
            .post(requestBody)
            .build();


        try {
            Response response = okHttpClient.newCall(request).execute();
            if (response.isSuccessful()) {
                log.debug(response.body().string());
                JSONObject jsonObject = JSONObject.parseObject(response.body().string());
                return "data:image/jpeg;base64," + jsonObject.getString("image");
            }
            else {
                throw new ServiceException("服务异常");
            }
        } catch (Exception e) {
            throw new ServiceException("服务异常");
        }
    }
}
